import argparse
import logging
from pathlib import Path
from pydantic import BaseModel
from pysam import VariantHeader, tabix_compress, tabix_index  # pylint: disable=E0611
import polars as pl

logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(name)s - %(levelname)s - %(message)s")
logger = logging.getLogger("Prepare File")


class PrepareRMSK(BaseModel):
    input_file: Path
    output_file: Path

    def header(self) -> str:
        header = VariantHeader()
        header.info.add(id="RMSK", number=1, type="String", description="RepeatMasker")
        return next(filter(lambda x: x.startswith("##INFO"), str(header).splitlines())).strip()

    def process_rmsk(self):
        records = pl.read_csv(self.input_file, separator="\t", has_header=False, columns=[5, 6, 7, 10], new_columns=["chrom", "start", "end", "rmsk"])
        records = records.filter(pl.col("chrom").str.len_chars().le(5)).sort(by=["chrom", "start", "end"])
        output_txt = self.output_file.with_suffix("") if self.output_file.suffix == ".gz" else self.output_file
        output_gz = output_txt.with_suffix(".txt.gz")
        with open(output_txt, "w") as writer:
            writer.write(self.header() + "\n")
            writer.write(records.write_csv(separator="\t", include_header=False))
        tabix_compress(str(output_txt), str(output_gz), force=True)
        tabix_index(str(output_gz), preset="bed", force=True)
        output_txt.unlink()

    def run(self):
        self.process_rmsk()


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Prepare dbNSFP SNV")
    parser.add_argument("--input_file", "-i", type=Path, required=True, help="input, RMSK file")
    parser.add_argument("--output_file", "-o", type=Path, required=True, help="output, RMSK file")
    parser.set_defaults(func=lambda args: PrepareRMSK.model_validate(vars(args)).run())
    args = parser.parse_args()
    args.func(args)
